import torch
import gradio as gr

title = "Is Python not Passion"
desc = "Anything that can go wrong will go wrong."

base_conf = 0.25
base_iou = 0.68

# 加载自定义YOLOv5模型
model = torch.hub.load("./", "custom", path="runs/train/exp3/weights/best.pt", source="local")

# 定义图像检测函数
def image_dect(image, conf, iou, name=""):
    model.conf = conf
    model.iou = iou
    results = model(image)
    # 在此处可以根据需要使用 image_name 进行进一步处理
    return results.render()[0]

# 示例数据
examples = [
    ['testData/1.jpg', 0.65, 0.7],
    ['testData/2.jpg', 0.65, 0.7]
]

# 创建Gradio界面
with gr.Blocks(css=".gr-title {text-align: center;}") as demo:
    gr.Markdown(f"# {title}", elem_classes="gr-title")
    gr.Markdown(desc)

    with gr.Row():  # 上传图片和检测结果在同一行
        with gr.Column():
            image_input = gr.Image(label="上传图像")
            conf_slider = gr.Slider(minimum=0, maximum=1, value=base_conf, label="置信度阈值")
            iou_slider = gr.Slider(minimum=0, maximum=1, value=base_iou, label="IoU阈值")
        with gr.Column():
            output_image = gr.Image(label="检测结果")

    # 当图像上传或滑块改变时，自动触发检测
    def auto_detect(*args):
        return image_dect(*args)

    image_input.change(auto_detect, inputs=[image_input, conf_slider, iou_slider], outputs=output_image)
    conf_slider.change(auto_detect, inputs=[image_input, conf_slider, iou_slider], outputs=output_image)
    iou_slider.change(auto_detect, inputs=[image_input, conf_slider, iou_slider], outputs=output_image)

    # 示例组件
    gr.Markdown("### 请选择你的英雄")
    examples_component = gr.Examples(examples=examples, inputs=[image_input, conf_slider, iou_slider], outputs=output_image)

demo.launch()
